#include <gflags/gflags.h>
#include <butil/logging.h>
#include <brpc/server.h>
#include <brpc/channel.h>
#include <json2pb/pb_to_json.h>
#include <algorithm>
#include <cstring>
#include <functional>
#include <string>
#include <vector>
#include <mutex>
#include <thread>
#include "kv.pb.h"
#include "timer.h"

DEFINE_string(attachment, "", "Carry this along with requests");
DEFINE_string(protocol, "baidu_std", "Protocol type. Defined in src/brpc/options.proto");
DEFINE_string(connection_type, "", "Connection type. Available values: single, pooled, short");
DEFINE_int32(timeout_ms, 100, "RPC timeout in milliseconds");
DEFINE_int32(port, 31011, "server brpc port");
DEFINE_string(listen_addr, "",
              "Server listen address, may be IPV4/IPV6/UDS."
              " If this is set, the flag port will be ignored");
DEFINE_int32(brpc_thread_num, 16, "server brpc thread_num");
DEFINE_string(brpc_service_name, "grpc_mock_dpu_server", "service name");
DEFINE_int32(idle_timeout_s, -1,
             "Connection will be closed if there is no "
             "read/write operations during the last `idle_timeout_s'");
DEFINE_string(load_balancer, "", "The algorithm for load balancing");

class EmbeddingTable {
public:
    int embedding_dim;
    std::string type;
    int num_embeddings;
    std::unordered_map<int, std::vector<float>> embeddings;

    std::mutex mtx;

    EmbeddingTable() {}

    EmbeddingTable(int dim, const std::string& t, int num)
        : embedding_dim(dim), type(t), num_embeddings(num) {
        std::vector<float> tmp_embedding;
        tmp_embedding.resize(dim);
        static int cnt = 0;
        // cnt += 1;
        cnt = 0;
        // std::fill(tmp_embedding.begin(), tmp_embedding.end(), 0.0f);
        std::fill(tmp_embedding.begin(), tmp_embedding.end(), (float)(cnt));
        for (int i = 0; i < num_embeddings; ++i) {
          /* just for easy debug */
          std::fill(tmp_embedding.begin(), tmp_embedding.end(), (float)(i));
          /* just for easy debug */
          set_embedding(i, tmp_embedding);
        }
    }

    void add_embedding(int id, std::vector<float> embedding) {
        if ((int)embedding.size() != embedding_dim) {
            std::cerr << "Error: Embedding size does not match embedding_dim\n";
            return;
        }
        std::lock_guard<std::mutex> lock(mtx);
        auto & emb = embeddings[id];
        for (int i = 0; i < embedding_dim; ++i) {
          emb[i] -= embedding[i];
        }
    }

    void set_embedding(int id, std::vector<float> embedding) {
        if ((int)embedding.size() != embedding_dim) {
            std::cerr << "Error: Embedding size does not match embedding_dim\n";
            return;
        }
        embeddings[id] = embedding;
    }

    std::vector<float> get_embedding(int id) const {
        auto it = embeddings.find(id);
        if (it != embeddings.end()) {
            return it->second;
        } else {
            std::cerr << "Error: Embedding not found\n";
            return {};
        }
    }
};

class KVStore {
public:
    std::unordered_map<std::string, std::shared_ptr<EmbeddingTable> > tables;

    static KVStore& getInstance() {
        static KVStore instance;
        return instance;
    }

    void create_table(const std::string& table_name, int embedding_dim, const std::string& type, int num_embeddings) {
        if (tables.find(table_name) != tables.end()) {
            std::cerr << "Error: Table already exists\n";
            return;
        }
        tables[table_name] = std::make_shared<EmbeddingTable>(embedding_dim, type, num_embeddings);
    }

    void add_embedding(const std::string& table_name, int id, const std::vector<float>& embedding) {
        auto it = tables.find(table_name);
        if (it != tables.end() && it->second != nullptr) {
            it->second->add_embedding(id, embedding);
        } else {
            std::cerr << "Error: Table not found\n";
        }
    }

    std::vector<float> get_embedding(const std::string& table_name, int id) const {
        auto it = tables.find(table_name);
        if (it != tables.end() && it->second != nullptr) {
            return it->second->get_embedding(id);
        } else {
            std::cerr << "Error: Table not found\n";
            return {};
        }
    }
};


namespace ks {
namespace reco {
class DocaDpuServiceImpl final : public DocaDpuService {
 public:
  DocaDpuServiceImpl() {

  }
  ~DocaDpuServiceImpl() {}
  void Initkvstore(::google::protobuf::RpcController *controller, const ::ks::reco::Initkv *request,
                  ::ks::reco::Initresponse *response, ::google::protobuf::Closure *done) override {
    (void) controller;
    brpc::ClosureGuard done_guard(done);
    KVStore& kvStore = KVStore::getInstance();
    int embedding_dim = request->embedding_dim();
    auto type = request->datatype();
    printf("Start Create kv_store\n");
    int table_size = request->table_names_size();
    for (int i = 0; i < table_size; ++i) {
      auto &table_name = request->table_names(i);
      auto num_embedding = request->num_embeddings(i);
      kvStore.create_table(table_name, embedding_dim, type, num_embedding);
    }
    /* the rpc status returned */
    printf ("Finish Create kv_store\n");
    response->set_status(1);

    return;
  }

  virtual void DpueFetchByTable(::google::protobuf::RpcController* controller,
                       const ::ks::reco::DpuFetchRequest* request,
                       ::ks::reco::DpuFetchResponse* response,
                       ::google::protobuf::Closure* done) {
    brpc::ClosureGuard done_guard(done);
    (void) controller;
    // std::thread::id this_id = std::this_thread::get_id();
    // std::cout << "Current thread ID: " << this_id << std::endl;

    xmh::Timer kv_store_timer("kvstore_server_time");

    KVStore& kvStore = KVStore::getInstance();

    int key_size = request->key_size();
    auto &table_name = request->table_name();
    auto option_type = request->option_type();
    // auto send_emb_size = request->send_embeddings_size();

    // auto recieve_debug = [&]() {
    //   printf ("key_info {key_size : %d}\n", key_size);
    //   printf ("table_name : %s\n", table_name.data());
    //   printf ("option_type : %s\n", option_type.data());
    //   printf ("fill_emb_size : %d\n", send_emb_size);
    // };
    // recieve_debug();

    for (int i = 0; i < key_size; ++i) {
      if (option_type == "get") {
        // printf ("option : get\n");
        std::vector<float> cur_emb = kvStore.get_embedding(table_name, request->key(i));
        int emb_size = (int)cur_emb.size() * sizeof(float);

        auto h_vector = response->add_embeddings();
        h_vector->resize(emb_size);
        std::memcpy(const_cast<char*>(h_vector->data()), cur_emb.data(), emb_size);
      }
      else if (option_type == "put") {
        // printf ("option : put\n");
        std::vector <float> decoder;
        auto & embedding_data = request->send_embeddings(i);
        int embedding_size = request->send_embeddings(i).size() / sizeof(float);

        for (int k = 0; k < embedding_size; ++k) {
            float value;
            memcpy(&value, embedding_data.data() + k * sizeof(float), sizeof(float));
            decoder.push_back(value);
        }
        kvStore.add_embedding(table_name, request->key(i), decoder);
        //TODO: change add_embedding type to (table_name, vector<int> keys, vector<vector<float>>)
      }
      else {
        printf ("invalid option found\n");
      }
    }

    kv_store_timer.end();

    return;
  }
 
 private:
  std::vector<std::vector<int>> table_ids = {
    {0, 1, 2, 4},
    {3},
    {5}
  };
};
}  // namespace reco
}  // namespace ks

int main() {
  // Parse gflags. We recommend you to use gflags as well.

  // Generally you only need one Server.
  LOG(INFO) << "server prepare to start";
  brpc::Server server;

  ks::reco::DocaDpuServiceImpl service_impl;
  if (server.AddService(&service_impl, brpc::SERVER_DOESNT_OWN_SERVICE) != 0) {
    LOG(ERROR) << "Fail to add service";
    return -1;
  }
  butil::EndPoint point;
  if (!FLAGS_listen_addr.empty()) {
    if (butil::str2endpoint(FLAGS_listen_addr.c_str(), &point) < 0) {
      LOG(ERROR) << "Invalid listen address:" << FLAGS_listen_addr;
      return -1;
    }
  } else {
    point = butil::EndPoint(butil::IP_ANY, FLAGS_port);
  }
  // Start the server.
  brpc::ServerOptions options;
  options.num_threads = 20;
  options.idle_timeout_sec = FLAGS_idle_timeout_s;

  xmh::Reporter::StartReportThread();

  if (server.Start(point, &options) != 0) {
    LOG(ERROR) << "Fail to start EchoServer";
    return -1;
  }

  // Wait until Ctrl-C is pressed, then Stop() and Join() the server.
  server.RunUntilAskedToQuit();

  xmh::Reporter::StopReportThread();
  xmh::Reporter::Report();
  /* close the connection */
  // wait_for_server();
  return 0;
}